package com.sunzm.spark.sql.udaf

import org.apache.spark.sql.{Encoder, Encoders}
import org.apache.spark.sql.expressions.Aggregator
import org.roaringbitmap.RoaringBitmap

/**
 * 把给定的一组Int类型的数据, 转化成 BitMap返回
 *
 * 注: 因为Spark没有提供 RoaringBitmap 类型的Encoder,自己实现比较麻烦, 所以考虑把 RoaringBitmap 序列化成 Array[Byte]
 *
 * RoaringBitmap不需要单独引入
 */
class IntToBitMapUDAF extends Aggregator[Int,Array[Byte],Array[Byte]]{
  /**
   * 初始化中间缓存结果
   * @return
   */
  override def zero: Array[Byte] = {
    // 构造一个空的bitmap
    val bm: RoaringBitmap = RoaringBitmap.bitmapOf()

    // 序列化这个空的bitmap交还给框架
    BitMapUtil.serBitMap(bm)
  }

  /**
   * 分区内聚合
   * @param buffer
   * @param element
   * @return
   */
  override def reduce(buffer: Array[Byte], element: Int): Array[Byte] = {

    // 将序列化后的buff反序列化成 RoaringBitMap对象
    val bitmap: RoaringBitmap = BitMapUtil.deSerBitMap(buffer)

    // 向缓存bitmap中添加新元素  element
    bitmap.add(element)

    // 再将更新后的缓存序列化成字节数组返回
    BitMapUtil.serBitMap(bitmap)

  }

  /**
   * 合并不同分区间的 buffer
   * @param buffer1
   * @param buffer2
   * @return
   */
  override def merge(buffer1: Array[Byte], buffer2: Array[Byte]): Array[Byte] = {
    BitMapUtil.mergeBitmapByteArray(buffer1, buffer2)
  }

  /**
   * 返回最终结果
   * @param reduction
   * @return
   */
  override def finish(reduction: Array[Byte]): Array[Byte] = {
    //这里直接返回是因为方法定义的时候,声明了返回值类型为 Array[Byte], 假设方法定义返回值类型为 Int, 还需要想办法转换成 Int 再返回
    reduction
  }

  /**
   * 输入数据的编码器
   * @return
   */
  override def bufferEncoder: Encoder[Array[Byte]] = Encoders.BINARY

  /**
   * 输出数据的编码器
   * @return
   */
  override def outputEncoder: Encoder[Array[Byte]] = Encoders.BINARY
}
